'''使用说明：
train_dir：train image 路径
validation_dir：validation image 路径
output_dir：tfrecord 输出路径
注意：此脚本需用到tensorflow下的object_detection.utils工具，所以需要将object_detection加到环境变量中'''


import io
import logging
import os
import random

import PIL.Image
import tensorflow as tf

from object_detection.utils import dataset_util

train_dir = 'your train image path'
validation_dir = 'your validation image path'
output_dir = 'tfrecord file output path_'

def dict_to_tf_example(image_subdirectory, filename):
  img_path = os.path.join(image_subdirectory, filename)
  #print(filename)
  with tf.gfile.GFile(img_path, 'rb') as fid:
    encoded_jpg = fid.read()
  encoded_jpg_io = io.BytesIO(encoded_jpg)
  image = PIL.Image.open(encoded_jpg_io)
  if image.format != 'JPEG':
    raise ValueError('Image format not JPEG')

  width = image.size[0]
  height = image.size[1]

  label = int(filename.split('_')[2].split('.')[0])
  feature_dict = {
      'image/height': dataset_util.int64_feature(height),
      'image/width': dataset_util.int64_feature(width),
      'image/encoded': dataset_util.bytes_feature(encoded_jpg),
      'image/format': dataset_util.bytes_feature('jpeg'.encode('utf8')),
      'image/class/label': dataset_util.int64_feature(label),
  }

  example = tf.train.Example(features=tf.train.Features(feature=feature_dict))
  return example


def create_tf_record(output_filename,
                     image_dir,
                     examples):
  writer = tf.python_io.TFRecordWriter(output_filename)
  for idx, example in enumerate(examples):
    if idx % 100 == 0:
      logging.info('On image %d of %d', idx, len(examples))
    try:
      #print(example)
      tf_example = dict_to_tf_example(image_dir, example)
      writer.write(tf_example.SerializeToString())
    except ValueError:
      logging.warning('Invalid example: %s, ignoring.', example)

  writer.close()


# TODO(derekjchow): Add test for pet/PASCAL main files.
def main(_):
  train_list = os.listdir(train_dir)
  validation_list = os.listdir(validation_dir)
  random.seed(42)
  random.shuffle(train_list)
  train_examples = len(train_list)
  validation_examples = len(validation_list)
  train_split_num = train_examples // 4 +1
  validation_split_num = validation_examples // 4 +1
  print(train_split_num)
  print(validation_split_num)
  for i in range(4):
    train_output_name = os.path.join(output_dir, ("pj_vehicle_train_0000{}-of-00004.tfrecord".format(i)))
    val_output_name = os.path.join(output_dir, ("pj_vehicle_validation_0000{}-of-00004.tfrecord".format(i)))
    create_tf_record(train_output_name, train_dir, train_list[train_split_num * i:train_split_num * (i + 1)])
    create_tf_record(val_output_name, validation_dir, validation_list[validation_split_num * i:validation_split_num * (i + 1)])

if __name__ == '__main__':
  tf.app.run()
